#!/usr/bin/python

from openpyxl.reader.excel import load_workbook

#获取混合样本的数据
def get_sample_data(filename):
    #保存样本数据的三维字典{样本号：{位点：{等位基因：峰高}}
    samples_data = {}
    f = open(filename,'r')
    while True:
        line = f.readline().rstrip()
        if not line:
            break
        lst = [a for a in line.split('\t') if a != '']
        if len(lst) == 4:
            sample_name = lst[0]
            if sample_name != 'Sample File Name':
                loci = lst[1]
                allele = lst[2]
                height = lst[3]
                if sample_name not in samples_data:
                    samples_data[sample_name] = {loci:{allele:height}}
                else:
                    if loci not in samples_data[sample_name]:
                        samples_data[sample_name][loci]= {allele:height}
                    else:
                        samples_data[sample_name][loci][allele] = height
    f.close()
    return samples_data

#读取男女样本的分型
def get_sample_genotype(filename):
    male_genotype = {}
    female_genotype = {}
    wb = load_workbook(filename,read_only=True)
    ws = wb.active
    for i in range(2,ws.max_row+1):
        loci = ws.cell(row=i, column=1).value
        male = str(ws.cell(row=i, column=3).value)
        female = str(ws.cell(row=i, column=2).value)
        male_alleles = [a for a in male.split(',') if a]
        female_alleles = [a for a in female.split(',') if a]
        male_genotype[loci] = male_alleles
        female_genotype[loci] = female_alleles
    wb.close()
    return male_genotype, female_genotype

#读取指定分型的样本的峰高
def get_sample_height(genotype,sample_data):
    sample_height = {}
    for loci in genotype:
        height_lst = []
        if loci in sample_data:
            for allele in genotype[loci]:
                if allele in sample_data[loci]:
                    height_lst.append(int(sample_data[loci][allele]))
        height = 0
        if len(height_lst) == 1:
            height = height_lst[0]/2
        if len(height_lst) == 2:
            height = height_lst[0]/2 + height_lst[1]/2
        sample_height[loci] = height
    return sample_height

#求单基因座比例
def get_ratio_loci(h1, h2):
    ratio = h1 / h2
    return ratio

def get_ratio_samples(samples_data,sample1st_genotype,sample2nd_genotype):
    ratio = []
    for sample in samples_data:
        sample1st_height = get_sample_height(sample1st_genotype, samples_data[sample])
        sample2nd_height = get_sample_height(sample2nd_genotype, samples_data[sample])
        for loci in sample2nd_height:
            h1 = sample1st_height[loci]
            h2 = sample2nd_height[loci]
            print(sample,loci,h1,h2)
            if sample2nd_height[loci] == 0:
                ratio.append([sample,loci, h1, h2,'无穷大'])
            else:
                ratio_value = '%.3f' % get_ratio_loci(h1, h2)
                ratio.append([sample,loci,h1, h2, ratio_value])
    return ratio